/*
 * Copyright The Slam Authors
 */

#include "slam/mapping/reflector_utilities.h"
#include <Eigen/Core>

namespace slam {
namespace mapping {
void ClusterPoints(const sensor::PointCloud& points, const float threshold,
                   sensor::PointCloud* centers_ptr,
                   std::vector<int>* count_ptr) {
  sensor::PointCloud& centers = *centers_ptr;
  std::vector<int>& count = *count_ptr;
  sensor::PointCloud center_sums;

  if (points.size() < 2u) {
    return;
  }

  centers.push_back(points[0]);
  center_sums.push_back(points[0]);
  count.push_back(1);

  for (int i = 1; i < static_cast<int>(points.size()); i++) {
    const Eigen::Vector3f& pt = points[i].position;
    int k = -1;
    for (int j = 0; j < static_cast<int>(centers.size()); j++) {
      const float dist = (pt - centers[j].position).norm();

      if (dist < threshold) {
        center_sums[j].position += pt;
        count[j]++;

        centers[j].position = (1.0 / count[j]) * center_sums[j].position;
        k = j;
        break;
      }
    }
    // Create new cluster center.
    if (k < 0) {
      centers.push_back(points[i]);
      center_sums.push_back(points[i]);
      count.push_back(1);
    }
  }
}

void ExtractReflectorFeature(const sensor::PointCloud& reflector_raw_data,
                             const int min_count,
                             sensor::PointCloud* reflector_data_ptr) {
  sensor::PointCloud& reflector_data = *reflector_data_ptr;

  sensor::PointCloud reflector_centers;
  std::vector<int> center_counts;
  ClusterPoints(reflector_raw_data, 0.1, &reflector_centers, &center_counts);

  reflector_data.clear();
  for (size_t i = 0u; i < center_counts.size(); ++i) {
#if 1
    const float x = std::fabs(reflector_centers[i].position(0));
    const float y = std::fabs(reflector_centers[i].position(1));
    const float theta = std::atan2(x, y);
    if (std::fabs(theta) < 35 / 180.0 * M_PI ||
        std::fabs(theta) > 75 / 180.0 * M_PI) {
      continue;
    }
#endif

    if (center_counts[i] >= min_count) {
      reflector_data.push_back(reflector_centers[i]);
    }
  }
}

void ReflectorMatching(const std::map<int, LandmarkPtr>& reflectors_map,
                       const transform::Rigid2d& T_local_tracking,
                       const sensor::PointCloud& reflector_feature_in_tracking,
                       const float match_distance_threshold,
                       std::vector<int>* match_reflector_ids_ptr,
                       std::vector<double>* match_min_dists_ptr) {
  if (reflector_feature_in_tracking.size() < 1) {
    return;
  }
  std::vector<int>& match_reflector_ids = *match_reflector_ids_ptr;
  std::vector<double>& match_min_dists = *match_min_dists_ptr;
  match_reflector_ids =
      std::vector<int>(reflector_feature_in_tracking.size(), -1);

  std::vector<Eigen::Vector2f> points_target;
  std::vector<Eigen::Vector2f> points_source;

  for (auto& point : reflector_feature_in_tracking) {
    auto reflector_in_local =
        T_local_tracking.cast<float>() * point.position.head<2>();
    points_source.push_back(reflector_in_local);
  }

  std::vector<int> landmark_ids;

  for (const auto& kv : reflectors_map) {
    const auto& reflector_id = kv.first;
    points_target.push_back(kv.second->p_local_landmark.cast<float>());
    landmark_ids.push_back(reflector_id);
  }

  std::vector<std::pair<int, int>> matches;

  for (size_t idx_source = 0u; idx_source < points_source.size();
       ++idx_source) {
    float min_dist = 100;
    int min_idx = -1;
    for (size_t idx_target = 0u; idx_target < points_target.size();
         ++idx_target) {
      const float dist =
          (points_target[idx_target] - points_source[idx_source]).norm();
      if (dist < min_dist) {
        min_dist = dist;
        min_idx = idx_target;
      }
    }
    if (min_idx != -1 && min_dist < match_distance_threshold) {
      matches.push_back(std::make_pair(min_idx, static_cast<int>(idx_source)));
    } else {
      LOG(INFO) << "No nearby reflector matched, min dist: " << min_dist;
    }
    match_min_dists.emplace_back(min_dist);
  }

  for (const auto& match : matches) {
    match_reflector_ids[match.second] = landmark_ids[match.first];
  }
}

void ReflectorMatching(
    const std::map<std::string, Eigen::Vector3d>& reflectors_nearby,
    const transform::Rigid3d& T_global_tracking,
    const sensor::PointCloud& reflector_feature_in_tracking,
    const float match_distance_threshold,
    std::vector<std::string>* match_reflector_ids_ptr,
    std::vector<double>* match_min_dists_ptr) {
  if (reflector_feature_in_tracking.size() < 1) {
    return;
  }
  std::vector<std::string>& match_reflector_ids = *match_reflector_ids_ptr;
  std::vector<double>& match_min_dists = *match_min_dists_ptr;
  match_reflector_ids =
      std::vector<std::string>(reflector_feature_in_tracking.size(), "");

  std::vector<Eigen::Vector2f> points_target;
  std::vector<Eigen::Vector2f> points_source;

  for (auto& point : reflector_feature_in_tracking) {
    auto reflector_in_global = T_global_tracking.cast<float>() * point;
    points_source.push_back(reflector_in_global.position.head<2>());
  }

  std::vector<std::string> landmark_ids;
  for (const auto& kv : reflectors_nearby) {
    const auto& reflector_id = kv.first;
    const Eigen::Vector3d& landmark_position = kv.second;
    points_target.push_back(landmark_position.cast<float>().head<2>());
    landmark_ids.push_back(reflector_id);
  }
  
  std::vector<std::pair<int, int>> matches;

  for (size_t idx_source = 0u; idx_source < points_source.size();
       ++idx_source) {
    float min_dist = 100;
    int min_idx = -1;
    for (size_t idx_target = 0u; idx_target < points_target.size();
         ++idx_target) {
      const float dist =
          (points_target[idx_target] - points_source[idx_source]).norm();
      if (dist < min_dist) {
        min_dist = dist;
        min_idx = idx_target;
      }
    }
    if (min_idx != -1 && min_dist < match_distance_threshold) {
      matches.push_back(std::make_pair(min_idx, static_cast<int>(idx_source)));
    } else {
      LOG(INFO) << "No nearby reflector matched, min dist: " << min_dist;
    }
    match_min_dists.emplace_back(min_dist);
  }

  for (const auto& match : matches) {
    match_reflector_ids[match.second] = landmark_ids[match.first];
  }
}

}  // namespace mapping
}  // namespace slam
